from abcframe import *
import model, numpy, ete2
## read http://pythonhosted.org/ete2/tutorial/tutorial_trees.html
## we might be able to be independent from this lib



class SimpleTree(Tree):
	_tree_str = ''
	_tree = None # tree structure, rooted tree for now
	_data = None # hmm this might be a dic. structure of arrays 
	_data_len = None # alignment length
	def __init__(self):
		self._tree_str = ''
		#super().__init__()
	def read_tree(self,f):
		try:
			fh = open(f,'r')
			# Todo: implement using ete2 library 
			fh.close()
		except: # take f as a string 
			self._tree_str = f			
			#self._tree = ete2.Tree(f) # parse a newick format
			self._tree = ete2.PhyloTree(f) # parse a newick format

	def read_data(self,f):
		self._tree.link_to_alignment(alignment=f, alg_format="fasta")
		for node in self._tree:
			if node.is_leaf():
				# convert character array to integer array for simple calculation
				d = map(int,list(node.sequence.replace('A','0').replace('C','1').replace('G','2').replace('T','2')))
				node.add_features(data=d)
				self._data_len = len(d)
		return 

	## maximum likelihood (felsenstein's pruning)  for now
	## for each position recursively calculate: 
    ## sum_y { Cj(y, children) * P(x,y ) }
	## Cj(X,v) = P(subtree with root v | vj = X)
	def _calc_llh(self, x, j, node, model):
		llh = 0.0
		t = node.dist
		if node.is_leaf(): # x is fixed and y is obtained from the data
			y = node.data[j] # alignment at j at species node.name
			llh = numpy.log(model.prob(x,y,t))
		else: # internal node, x is fixed  but sum up the probabilities of all y events
			for childnode in node.children:
				prob = 0 # not llh
				for y in range(4):
					## exp( log(x) + log(y)) = x * y
					prob += numpy.exp(numpy.log(model.prob(x,y,t)) +  self._calc_llh(y,j,childnode,model))
				llh += numpy.log(prob)
		return llh
		
	## tree traverse implement Felsenstein's pruning algorithm
	def calc_llh(self, model):
		root = self._tree.get_tree_root();
		llh = 0.0
		for j in range(self._data_len):
			for node in root.children: # traverse the tree
				for x in range(4):
					llh += self._calc_llh(x,j,node,model)
		return llh

	def __str__(self):
		return self._tree_str

			
if __name__ == '__main__':
	tree = SimpleTree()
	# newick tree
	tree.read_tree("(A:1,(B:1,(C:1,D:1):0.5):0.5);")
	# fasta alignments
	data = """
>A
AACG
>B
ACCG
>C
AACA
>D
AATG
"""
	tree.read_data(data)
	print tree._tree

	model = model.JC69Model()
	print tree.calc_llh(model)



